"""
批量标注脚本 - 处理整个数据集
支持断点续传，跳过已标注的图片
"""

import google.generativeai as genai
import os
import PIL.Image
import json
import cv2
import numpy as np
from pathlib import Path
from typing import List, Dict
import time
from annotate import CONFIG, AIAnalyzer, AnnotationGenerator, ImagePreprocessor

# 批量处理配置
BATCH_CONFIG = {
    "input_dir": "风机叶片数据集/images",  # 输入图片目录
    "output_dir": "风机叶片数据集/labels",  # 输出标注目录（会自动创建）
    "skip_existing": True,  # 是否跳过已标注的图片
    "save_visualization": True,  # 是否保存可视化图片
    "visualization_dir": "风机叶片数据集/visualizations",  # 可视化输出目录
    "delay_between_requests": 2,  # 请求间隔（秒），避免API限流
    "save_summary": True,  # 是否保存摘要报告
}

class BatchAnnotator:
    """批量标注器"""
    
    def __init__(self):
        self.input_dir = Path(BATCH_CONFIG["input_dir"])
        self.output_dir = Path(BATCH_CONFIG["output_dir"])
        self.viz_dir = Path(BATCH_CONFIG["visualization_dir"]) if BATCH_CONFIG["save_visualization"] else None
        
        # 创建输出目录
        self.output_dir.mkdir(parents=True, exist_ok=True)
        if self.viz_dir:
            self.viz_dir.mkdir(parents=True, exist_ok=True)
        
        # 初始化分析器
        genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
        self.analyzer = AIAnalyzer(model_name=CONFIG["ai_model_name"])
        
        # 统计信息
        self.stats = {
            "total": 0,
            "processed": 0,
            "skipped": 0,
            "failed": 0,
            "total_defects": 0,
            "defects_by_type": {cls: 0 for cls in CONFIG["defect_classes"]}
        }
    
    def get_image_list(self) -> List[Path]:
        """获取待处理的图片列表"""
        extensions = ["*.jpg", "*.jpeg", "*.png", "*.JPG", "*.JPEG", "*.PNG"]
        image_files = []
        
        for ext in extensions:
            image_files.extend(self.input_dir.glob(ext))
        
        # 如果启用跳过已标注，过滤掉已处理的
        if BATCH_CONFIG["skip_existing"]:
            filtered = []
            for img_path in image_files:
                label_path = self.output_dir / (img_path.stem + ".txt")
                if not label_path.exists():
                    filtered.append(img_path)
                else:
                    self.stats["skipped"] += 1
            
            print(f"📋 跳过已标注: {self.stats['skipped']} 张")
            return filtered
        
        return image_files
    
    def process_single_image(self, image_path: Path) -> bool:
        """处理单张图片"""
        try:
            print(f"\n{'='*60}")
            print(f"📷 处理: {image_path.name}")
            
            # 读取并预处理图片
            image_pil = PIL.Image.open(image_path)
            print(f"   尺寸: {image_pil.size[0]}x{image_pil.size[1]}")
            
            if CONFIG["enable_image_enhancement"]:
                image_pil = ImagePreprocessor.enhance_image(image_pil)
            
            # AI分析
            defects_data = self.analyzer.analyze_image(image_pil)
            
            if not defects_data:
                print("   ⚠️  未检测到缺陷")
                # 创建空标注文件
                label_path = self.output_dir / (image_path.stem + ".txt")
                label_path.touch()
                return True
            
            # 更新统计
            self.stats["total_defects"] += len(defects_data)
            for defect in defects_data:
                defect_type = defect.get("defect_type")
                if defect_type in self.stats["defects_by_type"]:
                    self.stats["defects_by_type"][defect_type] += 1
            
            print(f"   ✅ 检测到 {len(defects_data)} 个缺陷")
            
            # 生成标注
            generator = AnnotationGenerator(defects_data)
            
            # 保存YOLO格式标注到指定目录
            label_path = self.output_dir / (image_path.stem + ".txt")
            self._save_yolo_to_path(generator, label_path)
            
            # 保存可视化（如果启用）
            if BATCH_CONFIG["save_visualization"]:
                viz_path = self.viz_dir / (image_path.stem + "_annotated.jpg")
                self._save_visualization(generator, image_path, viz_path)
            
            self.stats["processed"] += 1
            return True
            
        except Exception as e:
            print(f"   ❌ 处理失败: {e}")
            self.stats["failed"] += 1
            return False
    
    def _save_yolo_to_path(self, generator: AnnotationGenerator, output_path: Path):
        """保存YOLO标注到指定路径"""
        annotations = []
        for defect in generator.defects:
            defect_type = defect.get("defect_type")
            box = defect.get("bounding_box")
            
            if defect_type in generator.class_to_index and box:
                class_index = generator.class_to_index[defect_type]
                yolo_line = generator._convert_to_yolo_format(box, class_index)
                annotations.append(yolo_line)
        
        if annotations:
            with open(output_path, 'w') as f:
                f.write("\n".join(annotations))
    
    def _save_visualization(self, generator: AnnotationGenerator, image_path: Path, output_path: Path):
        """保存可视化图片"""
        img_cv = cv2.imread(str(image_path))
        height, width, _ = img_cv.shape
        
        for defect in generator.defects:
            defect_type = defect.get("defect_type")
            box = defect.get("bounding_box")
            confidence = defect.get("confidence", 0)
            
            if defect_type in generator.class_to_color and box:
                color = generator.class_to_color[defect_type]
                x1 = int(box["x1"] * width)
                y1 = int(box["y1"] * height)
                x2 = int(box["x2"] * width)
                y2 = int(box["y2"] * height)
                
                # 绘制边界框
                cv2.rectangle(img_cv, (x1, y1), (x2, y2), color, 2)
                
                # 绘制标签
                label = f"{defect_type} {confidence:.2f}"
                cv2.putText(img_cv, label, (x1, y1 - 5),
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
        
        cv2.imwrite(str(output_path), img_cv)
    
    def run(self):
        """执行批量处理"""
        print("\n" + "="*60)
        print("🚀 批量标注系统启动")
        print("="*60)
        print(f"📁 输入目录: {self.input_dir}")
        print(f"📁 输出目录: {self.output_dir}")
        if self.viz_dir:
            print(f"📁 可视化目录: {self.viz_dir}")
        
        # 获取图片列表
        image_list = self.get_image_list()
        self.stats["total"] = len(image_list)
        
        if not image_list:
            print("\n⚠️  没有找到需要处理的图片")
            return
        
        print(f"\n📊 待处理: {len(image_list)} 张图片")
        print("="*60)
        
        # 处理每张图片
        start_time = time.time()
        
        for idx, image_path in enumerate(image_list, 1):
            print(f"\n进度: [{idx}/{len(image_list)}]")
            
            self.process_single_image(image_path)
            
            # API限流控制
            if idx < len(image_list):
                time.sleep(BATCH_CONFIG["delay_between_requests"])
        
        # 计算耗时
        elapsed_time = time.time() - start_time
        
        # 显示摘要
        self._print_summary(elapsed_time)
        
        # 保存摘要报告
        if BATCH_CONFIG["save_summary"]:
            self._save_summary_report(elapsed_time)
    
    def _print_summary(self, elapsed_time: float):
        """打印处理摘要"""
        print("\n" + "="*60)
        print("📊 处理摘要")
        print("="*60)
        print(f"总计图片: {self.stats['total']} 张")
        print(f"成功处理: {self.stats['processed']} 张")
        print(f"跳过已标注: {self.stats['skipped']} 张")
        print(f"处理失败: {self.stats['failed']} 张")
        print(f"总缺陷数: {self.stats['total_defects']} 个")
        print(f"\n缺陷类型分布:")
        for defect_type, count in self.stats["defects_by_type"].items():
            if count > 0:
                print(f"  - {defect_type}: {count} 个")
        print(f"\n⏱️  总耗时: {elapsed_time:.1f} 秒")
        print(f"⚡ 平均速度: {elapsed_time/max(self.stats['processed'], 1):.1f} 秒/张")
        print("="*60 + "\n")
    
    def _save_summary_report(self, elapsed_time: float):
        """保存摘要报告为JSON"""
        report = {
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "statistics": self.stats,
            "elapsed_time_seconds": elapsed_time,
            "config": {
                "model": CONFIG["ai_model_name"],
                "min_confidence": CONFIG["min_confidence"],
                "image_enhancement": CONFIG["enable_image_enhancement"],
            }
        }
        
        report_path = self.output_dir / "batch_summary.json"
        with open(report_path, 'w', encoding='utf-8') as f:
            json.dump(report, f, indent=2, ensure_ascii=False)
        
        print(f"📄 摘要报告已保存: {report_path}")


def main():
    """主函数"""
    try:
        annotator = BatchAnnotator()
        annotator.run()
    except KeyboardInterrupt:
        print("\n\n⚠️  用户中断，正在退出...")
    except Exception as e:
        print(f"\n❌ 批量处理出错: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()


